"""
https://blog.csdn.net/weixin_41424926/article/details/105383064

动手学习深度学习pytorch版——从零开始实现YOLOv1
"""
import os
import numpy as np
import cv2 as cv
from python_ai.CV_3.project.yolo.csdn_pt.prepare_data import GL_CLASSES, GL_NUMGRID, GL_NUMBBOX
from python_ai.common.xcommon import *

STATIC_DATASET_PATH = r'../../../../../../large_data/_very_large/VOCtrainval_11-May-2012/VOCdevkit/VOC2012_mini'
np.random.seed(1)

"""
.
├── Annotations                 source of annotations
├── ImageSets
│   ├── Action
│   ├── Layout
│   ├── Main
│   └── Segmentation
├── JPEGImages                  source of pictures
├── labels                      dest txt files format: cls and x,y,w,h
                                (1) Annotations => labels
                                (2) labels => labels for padding and scaling
├── SegmentationClass
├── SegmentationObject
└── voc2012_forYolov1           dest: csv and txt for train/test
                                csv format: (x,y,w,h,C)x2 + one_hotX20
    └── img                     dest: padded and scaled pictures

"""


def xywh2pts(x, y, w, h, idx_grids, w_input_size=448, h_input_size=448, w_grids=7, h_grids=7):
    gx = idx_grids % w_grids
    gy = idx_grids // w_grids
    w_grid_size = w_input_size / w_grids
    h_grid_size = h_input_size / h_grids

    base_x = gx * w_grid_size
    base_y = gy * h_grid_size
    x = base_x + w_grid_size * x
    y = base_y + h_grid_size * y
    w *= w_input_size
    h *= h_input_size

    x1 = x - w / 2
    y1 = y - h / 2
    x2 = x1 + w
    y2 = y1 + h

    x1 = int(x1)
    y1 = int(y1)
    x2 = int(x2)
    y2 = int(y2)
    return (x1, y1), (x2, y2)


sep('Load labels')
LABEL_DIR = os.path.join(STATIC_DATASET_PATH, 'voc2012_forYolov1')
txt_path = os.path.join(LABEL_DIR, 'train.txt')
csv_path = os.path.join(LABEL_DIR, 'train.csv')
print('Loading path ...')
with open(txt_path, 'r') as f:
    path_data = f.readlines()
print('Path loaded.')
print('Loading labels csv ...')
label_data = np.loadtxt(csv_path)
M, N = label_data.shape
print('label_data:', M, N)

sep('Check it')
offset = np.random.randint(0, M)
print('offset:', offset)
spr = 2
spc = 3
i = -1
result_img_arr = []
for row in range(spr):
    row_img_arr = []
    for col in range(spc):
        i += 1
        idx = i + offset
        idx %= M
        print('idx:', idx)

        img_path = '../' + path_data[idx]
        # img_path = img_path.replace('\\', '/')
        img_path = img_path.strip()
        print(f'path=|{img_path}|')
        img = cv.imread(img_path, cv.IMREAD_COLOR)

        label_row = label_data[idx]
        CELL_LEN = 30
        for j in range(GL_NUMGRID * GL_NUMGRID):
            label_cell = label_row[j * CELL_LEN:(j + 1) * CELL_LEN]
            if np.isclose(label_cell[4], 0):
                continue

            idx_cls = np.argmax(label_cell[10:])
            name_cls = GL_CLASSES[idx_cls]
            pt1, pt2 = xywh2pts(label_cell[0], label_cell[1], label_cell[2], label_cell[3], j)
            print(f'name_cls: {name_cls}, idx_cls: {idx_cls}, pt1: {pt1}, pt2: {pt2}')

            cv.rectangle(img, pt1, pt2, rand_color(), 2)
            cv.putText(img, name_cls, pt1, cv.FONT_HERSHEY_PLAIN, 2, rand_color(), 2)
        row_img_arr.append(img)
    row_img = np.concatenate(row_img_arr, axis=1)
    result_img_arr.append(row_img)
result_img = np.concatenate(result_img_arr, axis=0)
H, W = result_img.shape[:2]
if H / W > 9/16:
    rate = 800 / H
else:
    rate = 1600 / W
result_img = cv.resize(result_img, None, None, rate, rate, cv.INTER_CUBIC)
cv.imshow('Check', result_img)
cv.waitKey(0)
cv.destroyAllWindows()

sep('Over')
